import tensorflow as tf
from tools.dataset import Dataset


def broadcast_iou(box_1, box_2):
    # box_1: (..., (x1, y1, x2, y2))
    # box_2: (N, (x1, y1, x2, y2))

    # broadcast boxes
    box_1 = tf.expand_dims(box_1, -2)
    box_2 = tf.expand_dims(box_2, 0)
    # new_shape: (..., N, (x1, y1, x2, y2))
    new_shape = tf.broadcast_dynamic_shape(tf.shape(box_1), tf.shape(box_2))
    box_1 = tf.broadcast_to(box_1, new_shape)
    box_2 = tf.broadcast_to(box_2, new_shape)

    int_w = tf.maximum(tf.minimum(box_1[..., 2], box_2[..., 2]) -
                       tf.maximum(box_1[..., 0], box_2[..., 0]), 0)
    int_h = tf.maximum(tf.minimum(box_1[..., 3], box_2[..., 3]) -
                       tf.maximum(box_1[..., 1], box_2[..., 1]), 0)
    int_area = int_w * int_h
    box_1_area = (box_1[..., 2] - box_1[..., 0]) * \
                 (box_1[..., 3] - box_1[..., 1])
    box_2_area = (box_2[..., 2] - box_2[..., 0]) * \
                 (box_2[..., 3] - box_2[..., 1])
    return int_area / (box_1_area + box_2_area - int_area)


def yolo_boxes(pred, anchors, classes):
    # pred: (batch_size, grid_h, grid_w, anchors, (x, y, w, h, obj, ...classes))
    grid_size = tf.shape(pred)[1:3]
    box_xy, box_wh, objectness, class_probs = tf.split(
        pred, (2, 2, 1, classes), axis=-1)

    pred_box = tf.concat((box_xy, box_wh), axis=-1)  # original xywh for loss

    # !!! grid[x][y] == (y, x)
    grid = tf.meshgrid(tf.range(grid_size[1]), tf.range(grid_size[0]))
    grid = tf.expand_dims(tf.stack(grid, axis=-1), axis=2)  # [gx, gy, 1, 2]

    box_xy = tf.sigmoid(box_xy)
    box_xy = (box_xy + tf.cast(grid, tf.float32)) / \
             tf.cast([grid_size[1], grid_size[0]], tf.float32)
    box_wh = tf.exp(box_wh) * anchors

    box_x1y1 = box_xy - box_wh / 2
    box_x2y2 = box_xy + box_wh / 2
    bbox = tf.concat([box_x1y1, box_x2y2], axis=-1)

    return bbox, objectness, class_probs, pred_box


def YoloLoss(anchors, classes=80, batch_size=16, ignore_thresh=0.5,
             wh_weight=0.5, obj_weight=5.0, noobj_weight=0.5):
    def yolo_loss(y_true, y_pred):
        # 1. transform all pred outputs
        # y_pred: (batch_size, grid_h, grid_w, anchors, (x, y, w, h, obj, ...cls))
        pred_box, pred_obj, pred_class, pred_xywh = yolo_boxes(
            y_pred, anchors, classes)
        pred_xy = pred_xywh[..., 0:2]
        pred_wh = pred_xywh[..., 2:4]

        # 2. transform all true outputs
        # y_true: (batch_size, grid, grid, anchors, (x, y, w, h, obj, cls))
        true_box_, true_obj, true_class_idx = tf.split(
            y_true, (4, 1, 1), axis=-1)
        true_xy = true_box_[..., 0:2]
        true_wh = true_box_[..., 2:4]
        true_x1y1 = true_xy - true_wh / 2
        true_x2y2 = true_xy + true_wh / 2
        true_box = tf.concat([true_x1y1, true_x2y2], axis=-1)

        # give higher weights to small boxes
        box_loss_scale = 2 - true_wh[..., 0:1] * true_wh[..., 1:2]

        # 3. inverting the pred box equations
        grid_h = tf.shape(y_true)[1]
        grid_w = tf.shape(y_true)[2]
        grid = tf.meshgrid(tf.range(grid_w), tf.range(grid_h))
        grid = tf.expand_dims(tf.stack(grid, axis=-1), axis=2)

        grid_h = tf.expand_dims(grid_h, -1)
        grid_w = tf.expand_dims(grid_w, -1)
        grid_wh = tf.cast(tf.stack([grid_w, grid_h], axis=1), tf.float32)

        true_xy = true_xy * grid_wh - \
                  tf.cast(grid, tf.float32)
        true_wh = tf.math.log(true_wh / anchors)
        true_wh = tf.where(tf.math.is_inf(true_wh),
                           tf.zeros_like(true_wh), true_wh)

        # 4. calculate all masks
        obj_mask = tf.squeeze(true_obj, -1)
        # ignore false positive when iou is over threshold
        best_iou = tf.map_fn(
            lambda x: tf.reduce_max(broadcast_iou(x[0], tf.boolean_mask(
                x[1], tf.cast(x[2], tf.bool))), axis=-1),
            (pred_box, true_box, obj_mask),
            tf.float32)
        ignore_mask = tf.cast(best_iou < ignore_thresh, tf.float32)

        obj_mask = tf.expand_dims(obj_mask, axis=-1)
        ignore_mask = tf.expand_dims(ignore_mask, axis=-1)
        # 5. calculate all losses
        xy_loss = tf.reduce_sum(
            obj_mask * box_loss_scale * tf.nn.sigmoid_cross_entropy_with_logits(
                labels=true_xy, logits=pred_xy)) / batch_size

        wh_loss = tf.reduce_sum(
            obj_mask * box_loss_scale * wh_weight * tf.square(true_wh - pred_wh)) / batch_size

        obj_loss = obj_weight * tf.reduce_sum(
            obj_mask * tf.nn.sigmoid_cross_entropy_with_logits(
                labels=true_obj, logits=pred_obj)) / batch_size

        noobj_loss = noobj_weight * tf.reduce_sum(
            (1 - obj_mask) * ignore_mask * tf.nn.sigmoid_cross_entropy_with_logits(
                labels=true_obj, logits=pred_obj)) / batch_size

        true_class_idx = tf.squeeze(true_class_idx, axis=-1)
        class_loss = tf.reduce_sum(
            obj_mask * tf.nn.sigmoid_cross_entropy_with_logits(
                labels=tf.one_hot(tf.cast(true_class_idx, dtype=tf.int32), classes), logits=pred_class)) / batch_size

        return xy_loss + wh_loss + obj_loss + noobj_loss + class_loss
    return yolo_loss


if __name__ == "__main__":
    a = tf.constant([1, 2], dtype=tf.int32)
    b = tf.constant([3, 4], dtype=tf.int32)
    c = tf.concat([a, b], axis=-1)

    image_ann_path = 'data/k210_img_ann.npy'
    anchors_path = 'data/k210_anchor.npy'
    class_num = 10
    in_hw = (224, 320)
    batch_size = 64

    ds = Dataset(image_ann_path, anchors_path, batch_size, class_num, in_hw)
    anchors = ds.anchors
    anchor_masks = ds.anchor_masks

    loss = [YoloLoss(anchors[mask], classes=class_num)
            for mask in anchor_masks]

    y_pred = tf.ones(shape=(1, 7, 10, 3, 15), dtype=tf.float32)
    y_true = tf.ones(shape=(1, 7, 10, 3, 6), dtype=tf.float32)

    loss_ = loss[0](y_true, y_pred)